import torch.cuda


class FastLoader:
    def __init__(self, dataloader, device):
        super(FastLoader, self).__init__()
        self.q = []
        self.device = device
        self.dataloader = dataloader
        self.stream = torch.cuda.Stream()
        self.run()
        self.cnt = -1

    def run(self):
        with torch.cuda.stream(self.stream):
            for step, data in enumerate(self.dataloader, 0):
                inputs, labels = data[0].to(self.device), data[1].to(self.device)
                self.q.append((inputs, labels))

    def __next__(self):
        if self.cnt == len(self.q) - 1:
            raise StopIteration
        self.cnt = self.cnt + 1
        return self.q[self.cnt]

    def __iter__(self):
        self.cnt = -1
        return self
